from torch.utils.tensorboard import SummaryWriter
from PIL import Image
import numpy as np

if __name__ == "__main__":
    """
    SummaryWriter 可以保存训练过程中的数据。
    
    使用 tensorboard 命令可以浏览那些数据： tensorboard --logdir=demo02/logs --port=6006
    """

    writer = SummaryWriter("demo02/logs")

    writer.add_text("Hello World", "Hello")

    for i in range(10):
        writer.add_scalar("y=x", i+ 10, 3 * (i+ 3) +5 )

    image = Image.open("data/train/ants_image/6240329_72c01e663e.jpg")

    image_array = np.array(image)

    writer.add_image("Ant", image_array, 1, dataformats="HWC")

    image = Image.open("data/train/ants_image/7759525_1363d24e88.jpg")

    image_array = np.array(image)

    writer.add_image("Ant", image_array, 2, dataformats="HWC")

    writer.close()
